package com.wy.websocket.spring;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

import org.springframework.web.socket.WebSocketSession;

/**
 * WebSocket Session管理
 *
 * @author 飞花梦影
 * @date 2024-04-30 11:13:07
 * @git {@link https://github.com/dreamFlyingFlower}
 */
public class WebSocketSessionManager {

	/**
	 * 保存连接 session 的地方
	 */
	private static ConcurrentHashMap<String, WebSocketSession> SESSION_POOL = new ConcurrentHashMap<>();

	/**
	 * 添加 session
	 *
	 * @param key
	 */
	public static void add(String key, WebSocketSession session) {
		// 添加 session
		SESSION_POOL.put(key, session);
	}

	/**
	 * 删除 session,会返回删除的 session
	 *
	 * @param key
	 * @return
	 */
	public static WebSocketSession remove(String key) {
		// 删除 session
		return SESSION_POOL.remove(key);
	}

	/**
	 * 删除并同步关闭连接
	 *
	 * @param key
	 */
	public static void removeAndClose(String key) {
		WebSocketSession session = remove(key);
		if (session != null) {
			try {
				// 关闭连接
				session.close();
			} catch (IOException e) {
				// todo: 关闭出现异常处理
				e.printStackTrace();
			}
		}
	}

	/**
	 * 获得 session
	 *
	 * @param key
	 * @return
	 */
	public static WebSocketSession get(String key) {
		// 获得 session
		return SESSION_POOL.get(key);
	}
}